/**
 * ============================================================================
 *
 * Copyright (C) 2018, Hisilicon Technologies Co., Ltd. All Rights Reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *   1 Redistributions of source code must retain the above copyright notice,
 *     this list of conditions and the following disclaimer.
 *
 *   2 Redistributions in binary form must reproduce the above copyright notice,
 *     this list of conditions and the following disclaimer in the documentation
 *     and/or other materials provided with the distribution.
 *
 *   3 Neither the names of the copyright holders nor the names of the
 *   contributors may be used to endorse or promote products derived from this
 *   software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 * ============================================================================
 */
#include <sys/time.h>
#include <opencv2/opencv.hpp>
#include "dvpp_preprocess.h"
#include <vector>
#include "hiaiengine/log.h"
#include "hiaiengine/ai_types.h"
#include "hiaiengine/ai_model_parser.h"
#include "ascenddk/ascend_ezdvpp/dvpp_data_type.h"
#include "ascenddk/ascend_ezdvpp/dvpp_process.h"

using ascend::utils::DvppBasicVpcPara;
using ascend::utils::DvppProcess;
using ascend::utils::DvppVpcOutput;
using hiai::Engine;
using hiai::ImageData;
using hiai::IMAGEFORMAT;

namespace {
// output port (engine port begin with 0)
const uint32_t kSendDataPort = 0;

// face detection model input width
const uint32_t kModelWidth = 1080; // 1080

// face detection model input height
const uint32_t kModelHeight = 720; // 720

const uint32_t kModelOutputImageWidth = 540;
const uint32_t kModelOutputImageHeight = 360;

// call dvpp success
const uint32_t kDvppProcSuccess = 0;
// level for call DVPP
const int32_t kDvppToJpegLevel = 100;

// vpc input image offset
const uint32_t kImagePixelOffsetEven = 1;
const uint32_t kImagePixelOffsetOdd = 2;
}

// register custom data type
HIAI_REGISTER_DATA_TYPE("EngineTransT", EngineTransT);
HIAI_REGISTER_DATA_TYPE("OutputT", OutputT);
HIAI_REGISTER_DATA_TYPE("ScaleInfoT", ScaleInfoT);
HIAI_REGISTER_DATA_TYPE("NewImageParaT", NewImageParaT);
HIAI_REGISTER_DATA_TYPE("BatchImageParaWithScaleT", BatchImageParaWithScaleT);

DvppPreProcess::DvppPreProcess() {
  dataInputIn_ = nullptr;
  total_number_ = 0;
}

HIAI_StatusT DvppPreProcess::Init(
    const hiai::AIConfig& config,
    const std::vector<hiai::AIModelDescription>& model_desc) {
  HIAI_ENGINE_LOG("Start initialize!");
  
  return HIAI_OK;
}

HIAI_StatusT DvppPreProcess::ImagePreProcess(
    const ImageData<u_int8_t>& src_img, ImageData<u_int8_t>& resized_img) {
  if (src_img.format != IMAGEFORMAT::YUV420SP) {
    // input image must be yuv420sp nv12.
    HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT,
                    "[ODInferenceEngine] input image type does not match");
    return HIAI_ERROR;
  }

  // assemble resize param struct
  DvppBasicVpcPara dvpp_resize_param;
  dvpp_resize_param.input_image_type = INPUT_YUV420_SEMI_PLANNER_UV;
  dvpp_resize_param.src_resolution.height = src_img.height;
  dvpp_resize_param.src_resolution.width = src_img.width;

  // the value of crop_right and crop_left must be odd.
  dvpp_resize_param.crop_right =
      src_img.width % 2 == 0 ? src_img.width - kImagePixelOffsetEven :
      src_img.width-kImagePixelOffsetOdd;
  dvpp_resize_param.crop_down =
      src_img.height % 2 == 0 ? src_img.height - kImagePixelOffsetEven :
      src_img.height-kImagePixelOffsetOdd;

  dvpp_resize_param.dest_resolution.width = kModelWidth;
  dvpp_resize_param.dest_resolution.height = kModelHeight;

  // the input image is aligned in memory.
  dvpp_resize_param.is_input_align = true;

  DvppProcess dvpp_process(dvpp_resize_param);

  DvppVpcOutput dvpp_out;
  int ret = dvpp_process.DvppBasicVpcProc(src_img.data.get(),
      src_img.size, &dvpp_out);

  if (ret != kDvppProcSuccess) {
    HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT,
                    "call dvpp resize failed with code %d!",ret);
    return HIAI_ERROR;
  }

  // dvpp_out->pbuf
  resized_img.data.reset(dvpp_out.buffer, default_delete<uint8_t[]>());
  resized_img.size = dvpp_out.size;
  // HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT, "style_transfer src w:%d h:%d, dvpp output size:%d.", 
  //   src_img.width, src_img.height, dvpp_out.size);

  return HIAI_OK;
}


HIAI_IMPL_ENGINE_PROCESS("dvpp_preprocess",
    DvppPreProcess, INPUT_SIZE) {
  HIAI_ENGINE_LOG("Start process!");

  // check arg0 is null or not
  if (arg0 == nullptr) {
    HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT,
                    "Failed to process invalid message.");
    return HIAI_ERROR;
  }

  std::shared_ptr<BatchImageParaWithScaleT> image_handle =
      std::static_pointer_cast<BatchImageParaWithScaleT>(arg0);

  dataInputIn_ = std::make_shared<BatchImageParaWithScaleT>();

  dataInputIn_->b_info = image_handle->b_info;

  struct  timeval start;
  struct  timeval end;
  gettimeofday(&start, NULL);
  // initialize as zero
  uint32_t all_input_size = 0;
  // vector<ImageData<u_int8_t>> processed_imgs;
  for (uint32_t i = 0; i < image_handle->b_info.batch_size; i++) {
    // image pre process
    NewImageParaT image_param;
    HIAI_StatusT vpc_ret =
        ImagePreProcess(image_handle->v_img[i].img, image_param.img);
    if(vpc_ret != HIAI_OK){
      HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT,
                    "image pre process error");
      continue;
    }
    all_input_size += image_param.img.size * sizeof(uint8_t);
    dataInputIn_->v_img.push_back(image_param);
  }
  // input size is less than zero, do not need to inference
  if (all_input_size <= 0) {
    HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT,
                    "all input image size=%u is less than zero",
                    all_input_size);
    return HIAI_ERROR;
  }
  
  // send results and original image data to post process (port 0)
  HIAI_StatusT hiai_ret = SendData(kSendDataPort, "BatchImageParaWithScaleT",
                                   std::static_pointer_cast<void>(dataInputIn_));
  HIAI_ENGINE_LOG("End process!");
  return hiai_ret;
}
